# -*- coding: utf-8 -*-
"""
省份网格点提取模块

此模块提供从ERA5气象数据中提取指定省份的观测点经纬度数据的功能。
主要功能包括：
- 从ERA5数据中提取省份网格点
- 保存网格点数据为CSV文件
- 统计省份网格点数量
"""

import os
import pandas as pd
import numpy as np
import xarray as xr
from shapely.geometry import Point

from dry_wet_abrupt.config import PROVINCE_GRID_POINTS_DIR
from dry_wet_abrupt.gis.province_utils import read_province_shapefile, filter_provinces


class ProvinceGridPointExtractor:
    """
    省份网格点提取器
    用于从ERA5气象数据中提取指定省份的观测点经纬度数据
    """
    
    def __init__(self, shp_path, era5_path, name_field='NAME', output_dir=None):
        """
        初始化省份网格点提取器
        
        参数:
            shp_path (str): 省份边界GIS文件路径
            era5_path (str): ERA5数据文件路径
            name_field (str): 省份名称字段，默认为'NAME'
            output_dir (str): 输出目录，默认为配置中的省份网格点目录
        """
        self.shp_path = shp_path
        self.era5_path = era5_path
        self.name_field = name_field
        self.output_dir = output_dir or PROVINCE_GRID_POINTS_DIR
        
        # 确保输出目录存在
        os.makedirs(self.output_dir, exist_ok=True)
    
    def extract_and_save(self, provinces=None):
        """
        提取并保存省份网格点数据
        
        参数:
            provinces (list): 省份列表，如为None则提取所有省份
            
        返回:
            dict: 省份网格点数据字典
        """
        # 读取省份边界数据
        province_gdf = self._read_province_bounds()
        if province_gdf is None:
            return {}
        
        # 读取ERA5数据的经纬度信息
        lons, lats = self._read_era5_grid()
        if lons is None or lats is None:
            return {}
        
        # 提取指定省份的观测点经纬度数据
        province_data = self._extract_province_grid_points(province_gdf, lons, lats, provinces)
        if not province_data:
            return {}
        
        # 保存观测点数据为CSV文件
        self._save_to_csv(province_data)
        
        return province_data
    
    def _read_province_bounds(self):
        """
        从GIS文件中读取省份边界数据
        
        返回:
            geopandas.GeoDataFrame: 省份边界数据
        """
        try:
            # 读取shapefile文件
            gdf = read_province_shapefile(self.shp_path)
            print(f"成功读取省份边界GIS文件，共{len(gdf)}个省份")
            
            # 提取省份名称列表
            province_names = gdf[self.name_field].tolist()
            print(f"可用省份: {', '.join(province_names)}")
            
            return gdf
        except Exception as e:
            print(f"读取省份边界GIS文件失败: {e}")
            return None
    
    def _read_era5_grid(self):
        """
        读取ERA5数据的经纬度网格信息
        
        返回:
            tuple: (经度数组, 纬度数组)
        """
        try:
            ds = xr.open_dataset(self.era5_path)
            print(f"成功读取ERA5数据，维度: {ds.dims}")
            
            # 获取经纬度数组
            lons = ds.lon.values
            lats = ds.lat.values
            
            # 关闭数据集
            ds.close()
            
            print(f"经度范围: {lons.min():.2f} - {lons.max():.2f}, 分辨率: {np.diff(lons)[0]:.2f}")
            print(f"纬度范围: {lats.min():.2f} - {lats.max():.2f}, 分辨率: {np.diff(lats)[0]:.2f}")
            
            return lons, lats
        except Exception as e:
            print(f"读取ERA5数据失败: {e}")
            return None, None
    
    def _extract_province_grid_points(self, province_gdf, lons, lats, provinces=None):
        """
        提取指定省份的观测点经纬度数据
        
        参数:
            province_gdf (geopandas.GeoDataFrame): 省份边界数据
            lons (numpy.ndarray): 经度数组
            lats (numpy.ndarray): 纬度数组
            provinces (list): 要提取的省份名称列表，如为None则提取所有省份
            
        返回:
            dict: 包含各省份观测点经纬度数据的字典
        """
        # 如果未指定省份，则使用所有省份
        if provinces is None:
            provinces = province_gdf[self.name_field].tolist()
        else:
            # 筛选指定的省份
            filtered_gdf = filter_provinces(province_gdf, provinces, self.name_field)
            provinces = filtered_gdf[self.name_field].tolist()
        
        if not provinces:
            print("没有有效的省份可以提取")
            return {}
        
        # 创建结果字典
        result = {}
        
        # 对每个省份提取观测点
        for province in provinces:
            # 获取省份几何对象
            province_geom = province_gdf[province_gdf[self.name_field] == province]['geometry'].iloc[0]
            
            # 获取省份边界范围
            bounds = province_geom.bounds  # 返回(minx, miny, maxx, maxy)
            min_lon, min_lat, max_lon, max_lat = bounds
            
            # 提取在省份范围内的经纬度点
            province_lons = lons[(lons >= min_lon) & (lons <= max_lon)]
            province_lats = lats[(lats >= min_lat) & (lats <= max_lat)]
            
            # 创建经纬度网格点
            lon_grid, lat_grid = np.meshgrid(province_lons, province_lats)
            
            # 将网格点展平为一维数组
            grid_points_flat = np.column_stack((lon_grid.flatten(), lat_grid.flatten()))
            
            # 筛选真正在省份多边形内的点
            in_province_points = []
            for point in grid_points_flat:
                if province_geom.contains(Point(point[0], point[1])):
                    in_province_points.append(point)
            
            # 转换为numpy数组
            if in_province_points:
                grid_points = np.array(in_province_points)
            else:
                grid_points = np.empty((0, 2))
            
            # 保存到结果字典
            result[province] = {
                'grid_points': grid_points,
                'lon_count': len(province_lons),
                'lat_count': len(province_lats),
                'total_points': len(grid_points)
            }
            
            print(f"省份: {province}, 经度点数: {len(province_lons)}, 纬度点数: {len(province_lats)}, 有效点数: {len(grid_points)}")
        
        return result
    
    def _save_to_csv(self, province_data):
        """
        保存观测点数据为CSV文件
        
        参数:
            province_data (dict): 包含各省份观测点经纬度数据的字典
        """
        # 保存汇总信息
        summary_data = []
        for province, data in province_data.items():
            summary_data.append({
                '省份': province,
                '经度点数': data['lon_count'],
                '纬度点数': data['lat_count'],
                '有效点数': data['total_points']
            })
        
        summary_df = pd.DataFrame(summary_data)
        summary_file = os.path.join(self.output_dir, 'province_grid_points_summary.csv')
        summary_df.to_csv(summary_file, index=False, encoding='utf-8-sig')
        print(f"已保存汇总信息到: {summary_file}")
        
        # 为每个省份保存详细的观测点数据
        for province, data in province_data.items():
            grid_points = data['grid_points']
            if len(grid_points) > 0:
                df = pd.DataFrame(grid_points, columns=['经度', '纬度'])
                
                # 添加省份信息
                df['省份'] = province
                
                # 保存到CSV文件
                file_path = os.path.join(self.output_dir, f"{province}_grid_points.csv")
                df.to_csv(file_path, index=False, encoding='utf-8-sig')
                print(f"已保存{province}的观测点数据到: {file_path}")
            else:
                print(f"警告: {province}没有有效的观测点数据")
        
        # 保存所有省份的观测点数据到一个文件
        all_points = []
        for province, data in province_data.items():
            grid_points = data['grid_points']
            for point in grid_points:
                all_points.append({
                    '省份': province,
                    '经度': point[0],
                    '纬度': point[1]
                })
        
        if all_points:
            all_df = pd.DataFrame(all_points)
            all_file = os.path.join(self.output_dir, 'all_provinces_grid_points.csv')
            all_df.to_csv(all_file, index=False, encoding='utf-8-sig')
            print(f"已保存所有省份的观测点数据到: {all_file}")
        else:
            print("警告: 没有有效的观测点数据可以保存")